import os
import re
import sys
import json
import numpy as np
import locale
from datetime import datetime
from pathlib import Path
import folder_paths
import pprint
import piexif
import piexif.helper
import pymysql
import glob

# 导入原始SaveImageExtended的依赖
try:
    import pillow_avif
    avif_supported = True
except:
    avif_supported = False

try:
    from jxlpy import JXLImagePlugin
    jxl_supported = True
except:
    jxl_supported = False

from PIL import Image, ExifTags
from PIL.PngImagePlugin import PngInfo

version = "1.0.0"

class SaveImageExtendedXHH:
    """
    扩展的SaveImageExtended节点，支持LoRA信息提取和提示词管理
    使用MySQL数据库
    """
    
    RETURN_TYPES = ()
    FUNCTION = 'save_images'
    OUTPUT_NODE = True
    CATEGORY = 'ComfyUI-xhh/image'
    
    def __init__(self):
        self.output_dir = folder_paths.get_output_directory()
        self.prefix_append = ''
        
        # 图片质量设置
        self.png_compress_level = 9
        self.avif_quality = 60
        self.webp_quality = 75
        self.jpeg_quality = 91
        self.jxl_quality = 91
        self.tiff_quality = 91
        self.optimize_image = True
        
        # 文件名设置
        self.filename_prefix = 'ComfyUI'
        self.filename_keys = 'sampler_name, cfg, steps, %F %H-%M-%S'
        self.foldername_prefix = ''
        self.foldername_keys = 'ckpt_name'
        self.delimiter = '-'
        self.save_job_data = 'disabled'
        self.job_data_per_image = False
        self.job_custom_text = ''
        self.save_metadata = True
        self.counter_digits = 4
        self.counter_position = 'last'
        self.counter_positions = ['last', 'first']
        self.one_counter_per_folder = True
        self.image_preview = True
        self.extToRemove = ['.safetensors', '.ckpt', '.pt', '.bin', '.pth']
        self.output_ext = '.webp'
        self.output_exts = ['.webp', '.png', '.jpg', '.jpeg', '.gif', '.tiff', '.bmp']
        self.quality = 75
        
        # MySQL数据库配置
        self.mysql_config = {
            'host': 'localhost',
            'port': 3306,
            'user': 'root',
            'password': 'root',
            'database': 'comfyui-xhh',
            'charset': 'utf8mb4'
        }
        
        if jxl_supported:
            self.output_exts.insert(0, '.jxl')
        if avif_supported:
            self.output_exts.insert(0, '.avif')
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            'required': {
                'images': ('IMAGE', ),
                'filename_prefix': ('STRING', {'default': 'ComfyUI', 'multiline': False}),
                'filename_keys': ('STRING', {'default': 'sampler_name, cfg, steps, %F %H-%M-%S', 'multiline': True}),
                'foldername_prefix': ('STRING', {'default': '', 'multiline': False}),
                'foldername_keys': ('STRING', {'default': 'ckpt_name', 'multiline': True}),
                'delimiter': ('STRING', {'default': '-', 'multiline': False}),
                'save_job_data': ([
                    'disabled', 
                    'prompt', 
                    'basic, prompt', 
                    'basic, sampler, prompt', 
                    'basic, models, sampler, prompt'
                ], {'default': 'disabled'}),
                'job_data_per_image': ('BOOLEAN', {"default": False}),
                'job_custom_text': ('STRING', {'default': '', 'multiline': False}),
                'save_metadata': ('BOOLEAN', {'default': True}),
                'counter_digits': ('INT', {
                    "default": 4, 
                    "min": 1, 
                    "max": 8, 
                    "step": 1,
                    "display": "silder"
                }),
                'counter_position': (['last', 'first'], {'default': 'last'}),
                'one_counter_per_folder': ('BOOLEAN', {'default': True}),
                'image_preview': ('BOOLEAN', {'default': True}),
                'output_ext': (['.webp', '.png', '.jpg', '.jpeg', '.gif', '.tiff', '.bmp'], {'default': '.webp'}),
                'quality': ('INT', {
                    "default": 75, 
                    "min": 1, 
                    "max": 100, 
                    "step": 1,
                    "display": "silder"
                }),
                'extract_lora_info': ('BOOLEAN', {'default': True}),
                'save_to_prompt_manager': ('BOOLEAN', {'default': True}),
            },
            'optional': {
                'positive_text_opt': ('STRING', {'forceInput': True}),
                'negative_text_opt': ('STRING', {'forceInput': True}),
                'lora_names': ('STRING', {'forceInput': True}),
            },
            'hidden': {'prompt': 'PROMPT', 'extra_pnginfo': 'EXTRA_PNGINFO'},
        }
    
    def get_mysql_connection(self):
        """获取MySQL数据库连接"""
        try:
            connection = pymysql.connect(**self.mysql_config)
            return connection
        except Exception as e:
            print(f"[SaveImageExtendedXHH] MySQL连接失败: {e}")
            return None
    
    def generate_thumbnail(self, image_path, thumbnail_dir):
        """生成缩略图"""
        try:
            from PIL import Image
            
            # 确保缩略图目录存在
            os.makedirs(thumbnail_dir, exist_ok=True)
            
            # 打开原图
            with Image.open(image_path) as img:
                # 计算缩略图尺寸（保持宽高比）
                max_size = (512, 512)
                img.thumbnail(max_size, Image.Resampling.LANCZOS)
                
                # 生成缩略图文件名
                filename = os.path.basename(image_path)
                name, ext = os.path.splitext(filename)
                thumbnail_filename = f"{name}_thumb.jpg"
                thumbnail_path = os.path.join(thumbnail_dir, thumbnail_filename)
                
                # 保存缩略图
                img.save(thumbnail_path, "JPEG", quality=85)
                
                print(f"[SaveImageExtendedXHH] 生成缩略图: {thumbnail_path}")
                return thumbnail_path
                
        except Exception as e:
            print(f"[SaveImageExtendedXHH] 生成缩略图失败: {e}")
            return None
    
    def extract_lora_info_from_prompt(self, prompt, lora_names=None):
        """从工作流中提取LoRA信息"""
        lora_info = []
        
        # 处理手动输入的LoRA名称
        if lora_names and lora_names.strip():
            lora_name_list = [name.strip() for name in lora_names.split(',') if name.strip()]
            for lora_input in lora_name_list:
                # 判断输入的是路径还是名称
                if lora_input.endswith('.safetensors'):
                    # 如果是路径，提取文件名作为名称
                    lora_name = os.path.splitext(os.path.basename(lora_input))[0]
                    lora_path = lora_input
                else:
                    # 如果是名称，路径为空
                    lora_name = lora_input
                    lora_path = ''
                
                lora_info.append({
                    'name': lora_name,
                    'path': lora_path,
                    'node_id': 'manual_input',
                    'class_type': 'manual',
                    'source': 'manual'
                })
            print(f"[SaveImageExtendedXHH] 手动输入 {len(lora_name_list)} 个LoRA: {lora_name_list}")
        
        if not prompt:
            return lora_info
        
        try:
            for node_id, node_data in prompt.items():
                if isinstance(node_data, dict):
                    class_type = node_data.get('class_type', '')
                    inputs = node_data.get('inputs', {})
                    
                    # 检查是否是LoRA相关的节点
                    if 'lora' in class_type.lower() or 'lora' in str(inputs).lower():
                        lora_name = inputs.get('lora_name', '')
                        lora_path = inputs.get('lora_path', '')
                        
                        # 调试信息
                        print(f"[SaveImageExtendedXHH] 发现LoRA节点: {class_type}")
                        print(f"[SaveImageExtendedXHH] lora_name类型: {type(lora_name)}, 值: {lora_name}")
                        print(f"[SaveImageExtendedXHH] lora_path类型: {type(lora_path)}, 值: {lora_path}")
                        
                        # 处理非字符串类型的值
                        if isinstance(lora_name, (list, tuple)):
                            # 如果lora_name是列表，检查是否包含有效的LoRA信息
                            if lora_name and len(lora_name) > 0:
                                # 检查第一个元素是否是数字（可能是节点ID）
                                if isinstance(lora_name[0], (int, str)) and str(lora_name[0]).isdigit():
                                    print(f"[SaveImageExtendedXHH] 跳过数字ID: {lora_name[0]}")
                                    continue
                                lora_name = str(lora_name[0]) if lora_name else ''
                            else:
                                lora_name = ''
                        
                        if isinstance(lora_path, (list, tuple)):
                            lora_path = str(lora_path[0]) if lora_path else ''
                        
                        # 验证LoRA信息的有效性
                        if lora_name and not lora_name.isdigit() and len(lora_name) > 2:
                            lora_info.append({
                                'name': str(lora_name),
                                'path': str(lora_path),
                                'node_id': node_id,
                                'class_type': class_type
                            })
                        elif lora_path and lora_path.endswith('.safetensors'):
                            # 如果只有路径，从路径提取名称
                            lora_name = os.path.splitext(os.path.basename(lora_path))[0]
                            lora_info.append({
                                'name': lora_name,
                                'path': str(lora_path),
                                'node_id': node_id,
                                'class_type': class_type
                            })
                    
                    # 检查输入中是否包含LoRA路径
                    for key, value in inputs.items():
                        if isinstance(value, str) and ('lora' in key.lower() and value.endswith(('.safetensors'))):
                            lora_info.append({
                                'name': os.path.splitext(os.path.basename(value))[0],
                                'path': value,
                                'node_id': node_id,
                                'class_type': class_type,
                                'input_key': key
                            })
            
            print(f"[SaveImageExtendedXHH] 提取到 {len(lora_info)} 个LoRA信息")
            return lora_info
            
        except Exception as e:
            print(f"[SaveImageExtendedXHH] 提取LoRA信息失败: {e}")
            return []
    
    def save_lora_to_database(self, lora_info, image_path):
        """将LoRA信息保存到MySQL数据库"""
        try:
            conn = self.get_mysql_connection()
            if not conn:
                print(f"[SaveImageExtendedXHH] 无法连接到MySQL数据库")
                return
            
            with conn.cursor() as cursor:
                for lora in lora_info:
                    # 检查LoRA信息是否有效
                    if not lora.get('path') and not lora.get('name'):
                        print(f"[SaveImageExtendedXHH] 跳过无效的LoRA信息: {lora}")
                        continue
                    
                    # 检查是否已存在
                    print(f"[SaveImageExtendedXHH] 检查LoRA路径: {lora.get('path', 'N/A')}")
                    print(f"[SaveImageExtendedXHH] 检查LoRA名称: {lora.get('name', 'N/A')}")
                    
                    # 首先尝试按路径匹配
                    existing = None
                    if lora.get('path'):
                        cursor.execute("SELECT id FROM loras WHERE path = %s", (lora['path'],))
                        existing = cursor.fetchone()
                    
                    # 如果路径匹配失败，尝试按名称匹配
                    if not existing and lora.get('name'):
                        print(f"[SaveImageExtendedXHH] 尝试按名称匹配: {lora['name']}")
                        cursor.execute("SELECT id, path FROM loras WHERE name = %s", (lora['name'],))
                        existing = cursor.fetchone()
                        if existing:
                            print(f"[SaveImageExtendedXHH] 找到匹配的LoRA: {existing[1]}")
                    
                    # 如果名称匹配也失败，尝试按文件名匹配
                    if not existing and lora.get('path'):
                        filename = os.path.basename(lora['path'])
                        print(f"[SaveImageExtendedXHH] 尝试按文件名匹配: {filename}")
                        cursor.execute("SELECT id, path FROM loras WHERE path LIKE %s", (f'%{filename}',))
                        existing = cursor.fetchone()
                        if existing:
                            print(f"[SaveImageExtendedXHH] 找到匹配的LoRA: {existing[1]}")
                    
                    if existing:
                        # 获取找到的LoRA路径（可能是通过文件名匹配找到的）
                        if len(existing) > 1:
                            # 通过文件名匹配找到的，使用数据库中的完整路径
                            db_path = existing[1]
                            print(f"[SaveImageExtendedXHH] 使用数据库路径: {db_path}")
                        else:
                            # 精确匹配找到的，使用原始路径
                            db_path = lora['path']
                        
                        # 获取LoRA ID
                        lora_id = existing[0]
                        print(f"[SaveImageExtendedXHH] 找到现有LoRA: {lora_id}")
                        
                        # 使用绝对路径保存图片记录
                        filename = os.path.basename(image_path)
                        
                        # 生成缩略图
                        current_dir = os.path.dirname(os.path.abspath(__file__))
                        thumbnail_dir = os.path.join(current_dir, "..", "flux_prompt_helper", "promptManager-server", "files")
                        thumbnail_path = self.generate_thumbnail(image_path, thumbnail_dir)
                        
                        # 首先处理图片记录
                        cursor.execute("SELECT id FROM images WHERE filename = %s", (filename,))
                        image_record = cursor.fetchone()
                        
                        if not image_record:
                            # 创建新的图片记录
                            cursor.execute("""
                                INSERT INTO images (filename, file_path, thumbnail_path, created_at, updated_at)
                                VALUES (%s, %s, %s, NOW(), NOW())
                            """, (filename, image_path, thumbnail_path))
                            image_id = cursor.lastrowid
                            print(f"[SaveImageExtendedXHH] 创建新图片记录: {image_id}, 文件: {filename}")
                        else:
                            image_id = image_record[0]
                            print(f"[SaveImageExtendedXHH] 使用现有图片记录: {image_id}, 文件: {filename}")
                        
                        # 检查LoRA和图片的关联是否已存在
                        cursor.execute("SELECT id FROM lora_images WHERE lora_id = %s AND image_id = %s", 
                                     (lora_id, image_id))
                        existing_relation = cursor.fetchone()
                        
                        if not existing_relation:
                            # 创建新的关联
                            cursor.execute("""
                                INSERT INTO lora_images (lora_id, image_id, created_at)
                                VALUES (%s, %s, NOW())
                            """, (lora_id, image_id))
                            
                            # 更新LoRA时间戳
                            cursor.execute("""
                                UPDATE loras 
                                SET updated_at = NOW()
                                WHERE id = %s
                            """, (lora_id,))
                            
                            print(f"[SaveImageExtendedXHH] 为LoRA {lora_id} 添加图片关联: {image_id}")
                        else:
                            print(f"[SaveImageExtendedXHH] LoRA图片关联已存在，跳过: {image_id}")
                    else:
                        # LoRA不在数据库中，跳过添加新LoRA
                        print(f"[SaveImageExtendedXHH] LoRA不在数据库中，跳过: {lora['name']}")
                
                conn.commit()
            
        except Exception as e:
            print(f"[SaveImageExtendedXHH] 保存LoRA到数据库失败: {e}")
            import traceback
            traceback.print_exc()
        finally:
            if conn:
                conn.close()
    
    def save_prompt_to_manager(self, positive_text, negative_text, image_path):
        """将提示词保存到提示词管理器（多对一关系）"""
        print(f"[SaveImageExtendedXHH] 保存提示词到提示词管理器: positive_text: {positive_text}, negative_text: {negative_text}, image_path: {image_path}")
        try:
            conn = self.get_mysql_connection()
            if not conn:
                print(f"[SaveImageExtendedXHH] 无法连接到MySQL数据库")
                return
            
            with conn.cursor() as cursor:
                # 组合提示词
                full_prompt = f"Positive: {positive_text}\nNegative: {negative_text}" if negative_text else positive_text
                
                # 检查是否已存在相同的提示词
                cursor.execute("SELECT id FROM prompts WHERE text = %s", (full_prompt,))
                existing_prompt = cursor.fetchone()
                
                # 使用绝对路径保存图片记录
                filename = os.path.basename(image_path)
                
                # 生成缩略图
                current_dir = os.path.dirname(os.path.abspath(__file__))
                thumbnail_dir = os.path.join(current_dir, "..", "flux_prompt_helper", "promptManager-server", "files")
                thumbnail_path = self.generate_thumbnail(image_path, thumbnail_dir)
                
                # 首先处理图片记录
                cursor.execute("SELECT id FROM images WHERE filename = %s", (filename,))
                image_record = cursor.fetchone()
                
                if not image_record:
                    # 创建新的图片记录
                    cursor.execute("""
                        INSERT INTO images (filename, file_path, thumbnail_path, created_at, updated_at)
                        VALUES (%s, %s, %s, NOW(), NOW())
                    """, (filename, image_path, thumbnail_path))
                    image_id = cursor.lastrowid
                    print(f"[SaveImageExtendedXHH] 创建新图片记录: {image_id}, 文件: {filename}")
                else:
                    image_id = image_record[0]
                    print(f"[SaveImageExtendedXHH] 使用现有图片记录: {image_id}, 文件: {filename}")
                
                if existing_prompt:
                    # 更新现有提示词，添加新图片关联
                    prompt_id = existing_prompt[0]
                    
                    # 检查关联是否已经存在
                    cursor.execute("SELECT id FROM prompt_images WHERE prompt_id = %s AND image_id = %s", 
                                 (prompt_id, image_id))
                    existing_relation = cursor.fetchone()
                    
                    if not existing_relation:
                        # 添加新关联
                        cursor.execute("""
                            INSERT INTO prompt_images (prompt_id, image_id, created_at)
                            VALUES (%s, %s, NOW())
                        """, (prompt_id, image_id))
                        
                        # 更新提示词时间戳
                        cursor.execute("""
                            UPDATE prompts 
                            SET updated_at = NOW()
                            WHERE id = %s
                        """, (prompt_id,))
                        
                        print(f"[SaveImageExtendedXHH] 为提示词 {prompt_id} 添加图片关联: {image_id}")
                    else:
                        print(f"[SaveImageExtendedXHH] 图片关联已存在，跳过: {image_id}")
                else:
                    # 创建新提示词
                    import time
                    chrome_id = str(int(time.time() * 1000))  # 生成唯一的chrome_id
                    cursor.execute("""
                        INSERT INTO prompts (chrome_id, text, source, created_at, updated_at, last_synced)
                        VALUES (%s, %s, %s, NOW(), NOW(), NOW())
                    """, (chrome_id, full_prompt, 'workflow'))
                    
                    prompt_id = cursor.lastrowid
                    
                    # 添加图片关联
                    cursor.execute("""
                        INSERT INTO prompt_images (prompt_id, image_id, created_at)
                        VALUES (%s, %s, NOW())
                    """, (prompt_id, image_id))
                    
                    print(f"[SaveImageExtendedXHH] 创建新提示词 {prompt_id} 并添加图片关联: {image_id}")
                
                conn.commit()
                print(f"[SaveImageExtendedXHH] 提示词保存成功")
                
        except Exception as e:
            print(f"[SaveImageExtendedXHH] 保存提示词到管理器失败: {e}")
            import traceback
            traceback.print_exc()
        finally:
            if conn:
                conn.close()
    
    def cleanup_fileName(self, file='', extToRemove=None):
        """清理文件名"""
        if extToRemove is None:
            extToRemove = self.extToRemove
            
        if isinstance(file, str):
            for ext in extToRemove:
                file = file.removesuffix(ext)
        return file
    
    def find_keys_recursively(self, prompt={}, keys_to_find=[], found_values={}):
        """递归查找键值"""
        for key, value in prompt.items():
            if key in keys_to_find:
                found_values[key] = value
            elif isinstance(value, dict):
                self.find_keys_recursively(value, keys_to_find, found_values)
    
    def find_parameter_values(self, target_keys, prompt={}, found_values={}):
        """查找参数值"""
        loras_string = ''
        for key, value in prompt.items():
            if 'loras' in target_keys:
                if re.match(r'lora(_name)?(_\d+)?', key):
                    if value is not None:
                        value = self.cleanup_fileName(value)
                        loras_string += f'{value}, '
            
            if isinstance(value, dict):
                self.find_parameter_values(target_keys, value, found_values)
            
            if key in target_keys:
                value = self.cleanup_fileName(value)
                found_values[key] = value
        
        if 'loras' in target_keys and loras_string:
            found_values['loras'] = loras_string.strip().strip(',')
        
        if len(target_keys) == 1:
            return found_values.get(target_keys[0], None)
        
        return found_values
    
    def generate_custom_name(self, keys_to_extract, prefix, delimiter, prompt, timestamp=datetime.now()):
        """生成自定义文件名"""
        if '%' in prefix:
            custom_name = timestamp.strftime(prefix)
        else:
            custom_name = prefix
        
        if prompt is not None and keys_to_extract != ['']:
            found_values = {}
            
            for key in keys_to_extract:
                if not key:
                    continue
                
                value = None
                node, nodeKey = None, None
                
                if '/' in key and not '%' in key:
                    value = key
                else:
                    splitKey = key.split('.')
                    if len(splitKey) == 2 and not '%' in key:
                        if '' not in splitKey:
                            if splitKey[0].isdecimal():
                                node, nodeKey = splitKey[0], splitKey[1]
                                if node in prompt:
                                    self.find_keys_recursively(prompt[node], [nodeKey], found_values)
                                else:
                                    self.find_keys_recursively(prompt, [nodeKey], found_values)
                            else:
                                value = self.cleanup_fileName(key)
                        else:
                            value = key
                    else:
                        if '%' in key:
                            value = timestamp.strftime(key)
                        else:
                            nodeKey = key
                            self.find_keys_recursively(prompt, [nodeKey], found_values)
                
                if value is None:
                    if nodeKey is not None:
                        if nodeKey in found_values:
                            value = found_values[nodeKey]
                        if value is None:
                            value = nodeKey
                        else:
                            value = self.cleanup_fileName(value)
                
                if isinstance(value, str):
                    if custom_name.endswith('/'):
                        delim = ''
                    elif '/' in value and not value.endswith('/'):
                        delim = ''
                    elif value.startswith('.'):
                        delim = ''
                    else:
                        delim = delimiter
                else:
                    delim = delimiter
                
                if isinstance(value, float):
                    value = round(float(value), 1)
                
                custom_name += f"{delim}{value}"
        
        return custom_name.strip(delimiter).strip('.').strip('/').strip(delimiter)
    
    def get_latest_counter(self, one_counter_per_folder, folder_path, filename_prefix, counter_digits=4, counter_position='last', output_ext='.webp'):
        """获取最新计数器"""
        counter = 1
        if not os.path.exists(folder_path):
            return counter
        
        try:
            files = [file for file in os.listdir(folder_path) if file.endswith(output_ext)]
            extLen = len(output_ext)
            if files:
                if counter_position == 'last':
                    counters = [int(file[-(extLen + counter_digits):-extLen]) if file[-(extLen + counter_digits):-extLen].isdecimal() else 0 for file in files if one_counter_per_folder or file.startswith(filename_prefix)]
                else:
                    counters = [int(file[:counter_digits]) if file[:counter_digits].isdecimal() else 0 for file in files if one_counter_per_folder or file[counter_digits +1:].startswith(filename_prefix)]
                
                if counters:
                    counter = max(counters) + 1
        except Exception as e:
            print(f"[SaveImageExtendedXHH] 获取计数器失败: {e}")
        
        return counter
    
    def get_metadata_png(self, img, prompt, extra_pnginfo=None):
        """获取PNG元数据"""
        metadata = PngInfo()
        if prompt is not None:
            metadata.add_text('prompt', json.dumps(prompt))
        if extra_pnginfo is not None:
            for x in extra_pnginfo:
                metadata.add_text(x, json.dumps(extra_pnginfo[x]))
        return metadata
    
    def get_metadata_exif(self, img, prompt, extra_pnginfo=None):
        """获取EXIF元数据"""
        metadata = {}
        if prompt is not None:
            metadata["prompt"] = prompt
        if extra_pnginfo is not None:
            metadata.update(extra_pnginfo)
        
        exif = img.getexif()
        exif[0x010f] = "Prompt: " + json.dumps(metadata['prompt'])
        exif[0x010e] = "Workflow: " + json.dumps(metadata['workflow'])
        exif_dat = exif.tobytes()
        return exif_dat
    
    def save_image(self, image_path, img, prompt, save_metadata=True, extra_pnginfo=None, quality=75):
        """保存图片"""
        output_ext = os.path.splitext(os.path.basename(image_path))[1]
        metadata = None
        kwargs = dict()
        
        if output_ext in ['.avif']:
            if save_metadata:
                kwargs["exif"] = self.get_metadata_exif(img, prompt, extra_pnginfo)
            kwargs["quality"] = quality
            if quality == 100:
                kwargs["lossless"] = True
        elif output_ext in ['.webp']:
            if save_metadata:
                kwargs["exif"] = self.get_metadata_exif(img, prompt, extra_pnginfo)
            kwargs["quality"] = quality
            if quality == 100:
                kwargs["lossless"] = True
        elif output_ext in ['.jpg', '.jpeg']:
            if save_metadata:
                kwargs["exif"] = self.get_metadata_exif(img, prompt, extra_pnginfo)
            kwargs["quality"] = quality
            kwargs["optimize"] = self.optimize_image
        elif output_ext in ['.jxl']:
            if save_metadata:
                kwargs["exif"] = self.get_metadata_exif(img, prompt, extra_pnginfo)
            kwargs["quality"] = quality
            if quality == 100:
                kwargs["lossless"] = True
        elif output_ext in ['.tiff']:
            kwargs["quality"] = quality
            kwargs["optimize"] = self.optimize_image
        elif output_ext in ['.png', '.gif']:
            if save_metadata:
                kwargs["pnginfo"] = self.get_metadata_png(img, prompt, extra_pnginfo)
            kwargs["compress_level"] = self.png_compress_level
            kwargs["optimize"] = self.optimize_image
        
        img.save(image_path, **kwargs)
    
    def save_images(self, images, filename_prefix, filename_keys, foldername_prefix, foldername_keys, 
                   delimiter, save_job_data, job_data_per_image, job_custom_text, save_metadata, 
                   counter_digits, counter_position, one_counter_per_folder, image_preview, 
                   output_ext, quality, extract_lora_info, save_to_prompt_manager,
                   negative_text_opt=None, positive_text_opt=None, lora_names=None, extra_pnginfo=None, 
                   prompt=None):
        """保存图片的主方法"""
        
        # 应用默认值
        if not filename_prefix and not filename_keys:
            filename_prefix = self.filename_prefix
        if delimiter:
            delimiter = delimiter[0]
        
        filename_keys_to_extract = [item.strip() for item in filename_keys.split(',')]
        foldername_keys_to_extract = [item.strip() for item in foldername_keys.split(',')]
        
        timestamp = datetime.now()
        custom_filename = self.generate_custom_name(filename_keys_to_extract, filename_prefix, delimiter, prompt, timestamp)
        custom_foldername = self.generate_custom_name(foldername_keys_to_extract, foldername_prefix, delimiter, prompt, timestamp)
        
        # 获取分辨率
        i = 255. * images[0].cpu().numpy()
        img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
        resolution = f'{img.width}x{img.height}'
        
        # 提取LoRA信息
        lora_info = []
        if extract_lora_info:
            lora_info = self.extract_lora_info_from_prompt(prompt, lora_names)
            print(f"[SaveImageExtendedXHH] 提取到 {len(lora_info)} 个LoRA")
        
        # 创建文件夹并保存图片
        try:
            full_output_folder, filename, _, _, custom_filename = folder_paths.get_save_image_path(
                custom_filename, self.output_dir, images[0].shape[1], images[0].shape[0])
            output_path = os.path.join(full_output_folder, custom_foldername)
            os.makedirs(output_path, exist_ok=True)
            
            counter = self.get_latest_counter(one_counter_per_folder, output_path, filename, counter_digits, counter_position, output_ext)
            
            results = list()
            for image in images:
                i = 255. * image.cpu().numpy()
                img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
                
                if counter_position == 'last':
                    file = f'{filename}{delimiter}{counter:0{counter_digits}}{output_ext}'
                else:
                    file = f'{counter:0{counter_digits}}{delimiter}{filename}{output_ext}'
                
                image_path = os.path.join(output_path, file)
                self.save_image(image_path, img, prompt, save_metadata, extra_pnginfo, quality)
                
                # 保存LoRA信息到数据库
                if lora_info and save_to_prompt_manager:
                    self.save_lora_to_database(lora_info, image_path)
                
                # 保存提示词到管理器
                if save_to_prompt_manager and (positive_text_opt or negative_text_opt):
                    self.save_prompt_to_manager(positive_text_opt, negative_text_opt, image_path)
                
                results.append({'filename': file, 'subfolder': custom_foldername, 'type': 'output'})
                counter += 1
            
            if not image_preview:
                results = list()
            return {'ui': {'images': results}}
            
        except OSError as e:
            print(f"[SaveImageExtendedXHH] 保存图片失败: {e}")
            return {'ui': {'images': []}}

# 节点映射
NODE_CLASS_MAPPINGS = {
    'SaveImageExtendedXHH': SaveImageExtendedXHH,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    'SaveImageExtendedXHH': '💾 Save Image Extended XHH',
}
